# -*- coding: utf-8 -*-
# Purpose: Sample network
# Created: 12/08/2008
# Copyright (c) 2008,2009 by University College London
# Authors:
# Didrik Pinte <dpinte@dipole-consulting.com>
#
# This program is free software under the GPL (>=v2)
# Read the file COPYING coming with HydroPlatform for details.

import datetime
import numpy
import networkx

from hydroplatform.model import fields as fd
from hydroplatform.model import nodetypes as nd
from hydroplatform.model import database as db
from hydroplatform.model import network as nt
from hydroplatform.model import units
from numpy import array

def load_sample_network1(session):
    """
    Load a sample database into the given session
    """

    units.create_default_units(session)
    
    # create project
    project1 = db.Project(u'Project 1', 'Test project based on sample input file')
    project1.horizon_from = datetime.date(2000,1,1)
    project1.horizon_to = datetime.date(2002,12,1)
    project1.time_step = units.MONTHLY
    project1.projection_params = []
    project1.projection_name = u"My projection"
    project1.projection_epsg = None
    session.add(project1)
    session.commit()
    units.load_units(session, project1)
    
    # create network
    network1 = db.Network(u'Network 1', project1)
    session.add(network1)
    session.commit()
    
    # creating a node group
    groupp = db.FieldGroup(u"parameters")
    groupt = db.FieldGroup(u"tbl-phy")
    groupts = db.FieldGroup(u"time_series")
    groupd = db.FieldGroup(u"Default")
    session.add(groupp)
    session.add(groupt)
    session.add(groupts)
    session.commit()
    
    # create node types
    nodetypes = {
                 u'J': ('junction nodes', nd.SHP_CIRCLE, u'black', u'white', 5), 
                 u'SR': ('surface reservoir', nd.SHP_TRIANGLE, u'black', u'black', 5),
                 u'GW': ('groundwater reservoir', nd.SHP_SQUARE, u'blue', u'blue', 20),
                 u'Ag': ('agricultural demand nodes', nd.SHP_RECTANGLE, u'brown', u'brown', 20),
                 u'Urb': ('urban demand nodes', nd.SHP_ELLIPSE, u'black', u'black', 10),
                 u'Hp': ('hydropower plant nodes', nd.SHP_SQUARE, u'black', u'white', 10),                 
                 u'WWTP' : ('waste-water treatment plan nodes', nd.SHP_SQUARE, u'blue', u'white', 10),
                 u'Desal' : ('desalinisation nodes', nd.SHP_CIRCLE, u'blue', u'blue', 10)
    }    
    
    for (key, (desc, shape, pen, brush, size)) in nodetypes.items():
        nodetype = db.NodeType(key, desc, project1, shape, pen, brush, size)
        session.add(nodetype)
        nodetype.groups.append(groupd)
    session.commit()
    
    # set the node type attributes
    # J nodes
    jnodes = session.query(db.NodeType).filter_by(name=u"J").first()
    jnodes.groups.append(groupts)
    ts = fd.TimeSerie(        
        fromd = project1.horizon_from,
        tod   = project1.horizon_to,
        step  = project1.time_step,
        name="Boundary Condition",
        unit=u"m3/s")

    if jnodes is not None:
        j3_bc = db.NodeField(jnodes.id, ts, 
                             u"Boundary Condition", groupts)
        session.add(j3_bc)
        session.commit()
    else:
        raise Exception('Missing J node type in database')
    
    
    head_table = fd.Table(columns=[fd.TableCol(name="vol", unit=u"hm3"),
                                   fd.TableCol(name="surface", unit=u"hm3"),
                                   fd.TableCol(name="head", unit=u"m")],
                          name="tbl_phy")
    # SR nodes
    srnodes = session.query(db.NodeType).filter_by(name=u"SR").first()
    srnodes.groups.append(groupp)
    srnodes.groups.append(groupt)
    if srnodes is not None:
        maxvol        = db.NodeField(srnodes.id, fd.Parameter(unit=u"hm3",
                                                              name="MaxVol"), 
                                     u"MaxVol", groupp)
        maxrelease    = db.NodeField(srnodes.id, fd.Parameter(unit=u"m3/s",
                                                              name="MaxRelease"),
                                     u"MaxRelease", groupp)
        head_relation = db.NodeField(srnodes.id, 
                                    head_table, 
                                    u"tbl_phy", groupt)
        nref = db.NodeField(srnodes.id, fd.NodeReferences(name="nrefnode", data=[1]), u"nrefnode", groupp)
        lref = db.NodeField(srnodes.id, fd.LinkReferences(name="lrefnode", data=[1]), u"lrefnode", groupp)
        session.add(maxvol)
        session.add(maxrelease)
        session.add(head_relation)
        session.add(nref)
        session.add(lref)
        session.commit()
    else:
        raise Exception('Missing SR node type in database')
    
    nodes = [u'J1',u'J2',u'J3', u'SR1', u'SR2', u'SR3', u'SR4', u'GW1', u'GW2', \
             u'Ag1', u'Ag2', u'Urb1', u'Urb2', u'Hp1', u'WWTP1', u'Desal1']
    for node_name in nodes:
        # retrieve node type from node name 
        # FIXME : this is a really dirty way of doing that ...
        nodetype = session.query(db.NodeType).filter_by(name=node_name[:-1]).first()
        if nodetype is not None:
            node = db.Node(node_name, project1, nodetype)
        else :  
            raise Exception("Missing type for %s" % node_name)
        session.add(node)
    session.commit()
    
    # add attribute values
    # boundary condition for J3
    ts3data = numpy.transpose(array([[-15, -0.9, 35.2, 0.7,24.30,-16.00,   \
          -6.00,-6.80,-8.80,-11.00,-9.00,-3.00,2.20,19.10,54.00,10.00, \
          -20.00,-19.00,5.20,-13.00,-15.00,-19.00,-16.00,-5.00,-9.00,  \
          -12.00,-8.50,26.10,-1.40,-6.00,15.00,45.00,36.00,49.00,38.00,\
          23.20]]))
    node_j3 = session.query(db.Node).filter_by(name=u"J3").first()    
    nodevalue = db.NodeValue(node_j3.id, j3_bc.id, ts3data)
    session.add(nodevalue)    
    # parameters of SR nodes
    sr_parameters = {u'SR1' : (5000,100),
                     u'SR2' : (2000,75),
                     u'SR3' : (3000, 90),
                     u'SR4' : (6000,200)
                     }
    
    sr1tabledata = array([[100,10,1],
                 [500,50,5],
                 [1000,100,10],
                 [2000,200,20],
                 [5000,500,50]])
   
    mxv = maxvol.attribute_type.clone_traits()    
    mxr = maxrelease.attribute_type.clone_traits()
    for node_name, params in sr_parameters.items():
        node_sr = session.query(db.Node).filter_by(name=node_name).first()
        session.add(db.NodeValue(node_sr.id, maxvol.id, params[0]))
        session.add(db.NodeValue(node_sr.id, maxrelease.id, params[1]))
        if not node_name is u"SR1":
            session.add(db.NodeValue(node_sr.id, 
                                     head_relation.id, 
                                     sr1tabledata))
    

    linktype = db.LinkType(u"Default", u"No field associated", project1)
    session.add(linktype)
    linktype.size=1
    linktype.pen = u"black"
    
    links =[(u'J1',u'SR2'), (u'J2',u'J3'), (u'J2',u'Ag2'), (u'SR1',u'J1'),
            (u'SR1',u'Urb1'), (u'SR2',u'SR4'), (u'SR3',u'Hp1'), (u'SR4',u'J2'),
            (u'SR4',u'Ag1'), (u'SR4',u'Urb2'), (u'GW1',u'Urb1'),
            (u'GW2',u'Ag1'), (u'GW2',u'Ag2'),(u'Ag1',u'J2'), (u'Ag1',u'GW2'),
            (u'Ag2',u'J3'), (u'Ag2', u'GW2'), (u'Urb1',u'J1'), 
            (u'Urb1',u'GW1'), (u'Urb2',u'GW1'), (u'Urb2',u'WWTP1'),
            (u'Hp1',u'SR4'), (u'WWTP1', u'J2'), (u'WWTP1', u'Urb2'), 
            (u'Desal1',u'Urb2')]
    
    for (key,value) in links:
        query = session.query(db.Node)
        start_node = query.filter_by(name=key).first()
        end_node = query.filter_by(name=value).first()        
        link = db.Link(project1, start_node, end_node, linktype)
        session.add(link)
    session.commit
    
    network = session.query(db.Network).first()
    
    # add all the links and nodes to network 1
    for node in session.query(db.Node).all():
        network.nodes.append(node)
    session.commit()
    for link in session.query(db.Link).all():
        network.links.append(link)
    session.commit() 
    
    network.load_graph()
    pos = networkx.layout.spring_layout(network.graph)
    for node in network.nodes:
        coords = pos[node]
        node.x = coords[0]
        node.y = coords[1]
    session.commit()
    
    return (nodetypes, nodes, links, network)